# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Created by Tianheng Cheng(tianhengcheng@gmail.com)
# ------------------------------------------------------------------------------
import codecs
import csv
import os
import pprint
import argparse

import cv2
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
import lib.models as models
from lib.config import config, update_config
from lib.utils import utils
from lib.datasets import get_dataset
from lib.core import function


def parse_args():
    parser = argparse.ArgumentParser(description='Train Face Alignment')

    parser.add_argument('--cfg', help='experiment configuration filename',
                        required=True, type=str)
    parser.add_argument('--model-file', help='model parameters', required=True, type=str)

    args = parser.parse_args()
    update_config(config, args)
    return args


def main():
    args = parse_args()

    logger, final_output_dir, tb_log_dir = \
        utils.create_logger(config, args.cfg, 'test')

    logger.info(pprint.pformat(args))
    logger.info(pprint.pformat(config))

    cudnn.benchmark = config.CUDNN.BENCHMARK
    cudnn.determinstic = config.CUDNN.DETERMINISTIC
    cudnn.enabled = config.CUDNN.ENABLED

    config.defrost()
    config.MODEL.INIT_WEIGHTS = False
    config.freeze()
    model = models.get_face_alignment_net(config)

    # copy model files
    writer_dict = {
        'writer': SummaryWriter(log_dir=tb_log_dir),
        'train_global_steps': 0,
        'valid_global_steps': 0,
    }

    gpus = list(config.GPUS)
    # model = nn.DataParallel(model, device_ids=gpus).cuda()
    model.to(device=torch.device("cpu"))

    # load model
    # state_dict = torch.load(args.model_file)
    state_dict = torch.load(args.model_file, map_location="cpu")
    if 'state_dict' in state_dict.keys():
        state_dict = state_dict['state_dict']
        if os.path.islink(args.model_file):
            state_dict = state_dict.state_dict()
        model.module.load_state_dict(state_dict)
        # model.load_state_dict(state_dict.module.state_dict())
    else:
        model.module.load_state_dict(state_dict)

    dataset_type = get_dataset(config)

    predict_loader = DataLoader(
        dataset=dataset_type(config, is_train=False, is_predict=True),
        batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus),
        shuffle=False,
        num_workers=config.WORKERS,
        pin_memory=config.PIN_MEMORY
    )

    predictions, filenames = function.predict(config, predict_loader, model, writer_dict)

    torch.save(predictions, os.path.join(final_output_dir, 'predictions_all.pth'))

    utils.save_csv(predictions, filenames, save_path=os.path.join(final_output_dir, 'predictions_all.csv'))


if __name__ == '__main__':
    main()
